#include "cdsl.h"

#define color(n)                                \
  (!(n) ? BLACK : (n)->color)

#define copy_color(n,s)                         \
  ((n)->color = color(s))

#define set_black(n)                            \
  ((n)->color = BLACK)

#define set_red(n)                              \
  ((n)->color = RED)

#define is_black(n)                             \
  (color(n) == BLACK)

#define is_red(n)                               \
  (color(n) == RED)

#define rb_parent(n)                            \
  (tree_entry(parent(&(n)->node), struct rb_tree_node, node))

#define rb_set_parent(n, p)                                             \
  ((p) ? set_parent(&(n)->node, &(p)->node) : set_parent(&(n)->node, NULL))

#define rb_left(n)                              \
  (tree_entry(left(&(n)->node), struct rb_tree_node, node))

#define rb_set_left(p, c)                       \
  ((c) ? set_left(&(p)->node, &(c)->node) : set_left(&(p)->node, NULL))

#define rb_right(n)                             \
  (tree_entry(right(&(n)->node), struct rb_tree_node, node))

#define rb_set_right(p, c)                      \
  ((c) ? set_right(&(p)->node, &(c)->node) : set_right(&(p)->node, NULL))

#define rb_is_leaf(n)                           \
  (!(n))

#define rb_is_root(n)                           \
  (is_root(&(n)->node))

#define rb_is_child_left(n)                     \
  (is_child_left(&(n)->node))

#define rb_is_child_right(n)                    \
  (is_child_right(&(n)->node))

#define rb_grandparent(n)                                               \
  (tree_entry(grandparent(&(n)->node), struct rb_tree_node, node))

#define rb_uncle(n)                                             \
  (tree_entry(uncle(&(n)->node), struct rb_tree_node, node))

#define rb_sibling(n)                                           \
  (tree_entry(sibling(&(n)->node), struct rb_tree_node, node))

#define INIT_RB_TREE_NODE(n) do {                           \
      INIT_TREE_NODE(&(n)->node);                           \
      INIT_LIST_HEAD(&(n)->list);                           \
      set_red(n);                                           \
  } while(0)

#define RB_TREE_NODE_INIT(n)                                    \
  { RED, TREE_NODE_INIT(n.node), LIST_HEAD_INIT(n.list) }

#define rb_tree_free(t, n)                              \
  do { if ((t)->_free) (t)->_free(n); } while(0)

#define rb_tree_comp(t, n1, n2)                         \
  ((t)->_comp(n1, n2))

static void rotate_right(struct rb_tree *tree, struct rb_tree_node *node)
{
    struct rb_tree_node *child  = rb_left(node),
                        *parent = rb_parent(node);

    if (parent) {
        if (rb_is_child_left(node)) {
            rb_set_left(parent, child);
        }
        else {
            rb_set_right(parent, child);
        }
    }
    else {
        tree->rn = child;
    }

    rb_set_parent(child, parent);
    rb_set_parent(node, child);
    rb_set_left(node, rb_right(child));
    rb_set_right(child, node);

    if (!rb_is_leaf(rb_left(node))) {
        rb_set_parent(rb_left(node), node);
    }
}

static void rotate_left(struct rb_tree *tree, struct rb_tree_node *node)
{
    struct rb_tree_node *child  = rb_right(node),
                        *parent = rb_parent(node);

    if (parent) {
        if (rb_is_child_left(node)) {
            rb_set_left(parent, child);
        }
        else {
            rb_set_right(parent, child);
        }
    }
    else {
        tree->rn = child;
    }

    rb_set_parent(child, parent);
    rb_set_parent(node, child);
    rb_set_right(node, rb_left(child));
    rb_set_left(child, node);

    if (!rb_is_leaf(rb_right(node))) {
        rb_set_parent(rb_right(node), node);
    }
}

static void update_add(struct rb_tree *tree, struct rb_tree_node *node)
{
    int break_loop = FALSE;

    while (!break_loop) {
        break_loop = TRUE;

        if (rb_is_root(node)) {
            /* Case 1: this is the root node and it must be black 
             * Solution: repaint black
             */
            set_black(node);
        }
        else if (is_black(rb_parent(node))) {
            /* Case 2: the parent is black and so nothing has to be changed 
             */
        }
        else if (is_red(rb_uncle(node))) {
            /* Case 3: the parent and uncle are red (violate prop 5) 
             * Solution: repaint them as black and repaint the grandparent red
             */
            set_black(rb_parent(node));
            set_black(rb_uncle(node));
            node = rb_grandparent(node);
            set_red(node);
            break_loop = FALSE;
        }
        else {
            /* Case 4: parent red and uncle black
             * Solution: in some cases change the roles of parent and new node
             *     and go to case 5
             */
            if (rb_is_child_right(node) && rb_is_child_left(rb_parent(node))) {
                rotate_left(tree, rb_parent(node));
                node = rb_left(node);
            }
            else if (rb_is_child_left(node) && rb_is_child_right(rb_parent(node))) {
                rotate_right(tree, rb_parent(node));
                node = rb_right(node);
            }

            /* Case 5: parent red and uncle black, but with correct child place
             * Solution: preform appropriate rotation on grandparent after changing colors
             */
            set_black(rb_parent(node));
            set_red(rb_grandparent(node));
            if (rb_is_child_left(node)) {
                rotate_right(tree, rb_grandparent(node));
            }
            else {
                rotate_left(tree, rb_grandparent(node));
            }
        }
    }
}

void rb_tree_add(struct rb_tree *tree, struct rb_tree_node *val)
{
    INIT_RB_TREE_NODE(val);
    
    if (tree->rn) {
        struct rb_tree_node *rn = tree->rn;
        
        /* Find position and insert the new node */
        while (!rb_is_leaf(rn)) {
            rb_set_parent(val, rn);
            rn = (rb_tree_comp(tree, rn, val) > 0) ? rb_left(rn) : rb_right(rn);
        }
        if (rb_tree_comp(tree, rb_parent(val), val) > 0) {
            rb_set_left(rb_parent(val), val);
        }
        else {
            rb_set_right(rb_parent(val), val);
        }
    }
    else {
        /* Insert into the tree as root */
        tree->rn = val;
    }

    /* Insert into the sorted list */
    if (list_empty(&tree->fn)) {
        list_add_front(&val->list, &tree->fn);
    }
    else if (rb_is_child_left(val)) {
        list_add_back(&val->list, &rb_parent(val)->list);
    }
    else {
        list_add_front(&val->list, &rb_parent(val)->list);
    }

    update_add(tree, val);
}

static void rb_xchg_node(struct rb_tree *tree, struct rb_tree_node *n1, struct rb_tree_node *n2)
{
    xchg_node(&n1->node, &n2->node);

    /* Change the root node is necessary */
    if (tree->rn == n1) {
        tree->rn = n2; 
    }
    else if (tree->rn == n2) {
        tree->rn = n1;
    }
    
    /* Exchange colors to preserv tree properties */
    xchg(n1->color, n2->color);
}

static void do_balance(struct rb_tree *tree, struct rb_tree_node *node)
{
    int break_loop = FALSE;

    /* Remove the node from the rb tree */
    while (!break_loop) {
        break_loop = TRUE;

        if (rb_is_root(node)) {
            /* Case 1: the node to delete is the root
             * Solution: we are done 
             */
        }
        else {
            if (is_red(rb_sibling(node))) {
                /* Case 2: the sibling is red
                 * Solution: reverse sibling and parent and rotate
                 */
                set_red(rb_parent(node));
                set_black(rb_sibling(node));
                if (rb_is_child_left(node)) {
                    rotate_left(tree, rb_parent(node));
                }
                else {
                    rotate_right(tree, rb_parent(node));
                }
            }

            if (is_black(rb_sibling(node))       && 
                is_black(rb_parent(node))        &&
                is_black(rb_left(rb_sibling(node))) &&
                is_black(rb_right(rb_sibling(node)))) {
                /* Case 3: parent, sibling and children are all black
                 * Solution: repaint sibling red and goto case 1 for the parent
                 */
                set_red(rb_sibling(node));
                node = rb_parent(node);
                break_loop = FALSE;
            }
            else if (is_red(rb_parent(node))            && 
                     is_black(rb_sibling(node))         &&
                     is_black(rb_left(rb_sibling(node)))   &&
                     is_black(rb_right(rb_sibling(node)))) {
                /* Case 4: black node, red parent and black sibling with black children
                 * Solution: paint sibling red and parent black
                 */
                set_red(rb_sibling(node));
                set_black(rb_parent(node));
            }
            else {
                /* Case 5: complex stuff
                 * Solution: prepare tree for case 6 (see wiki online)
                 */
                if (is_black(rb_sibling(node))) {
                    if (rb_is_child_left(node)         &&
                        is_black(rb_right(rb_sibling(node))) &&
                        is_red(rb_left(rb_sibling(node)))) {
                        set_red(rb_sibling(node));
                        set_black(rb_left(rb_sibling(node)));
                        rotate_right(tree, rb_sibling(node));
                    }
                    else if (rb_is_child_right(node) &&
                             is_black(rb_left(rb_sibling(node))) &&
                             is_red(rb_right(rb_sibling(node)))) {
                        set_red(rb_sibling(node));
                        set_black(rb_right(rb_sibling(node)));
                        rotate_left(tree, rb_sibling(node));
                    }
                }
                
                /* Case 6: complex stuff
                 * Solution: final transformation to balance the tree
                 */
                copy_color(rb_sibling(node), rb_parent(node));
                set_black(rb_parent(node));
                if (rb_is_child_left(node)) {
                    set_black(rb_right(rb_sibling(node)));
                    rotate_left(tree, rb_parent(node));
                }
                else {
                    set_black(rb_left(rb_sibling(node)));
                    rotate_right(tree, rb_parent(node));
                }
            }
        }
    }
}

static void update_delete(struct rb_tree *tree, struct rb_tree_node *val)
{
    struct rb_tree_node *cn = 
      rb_is_leaf(rb_left(val)) ? rb_right(val) : rb_left(val);

    if (rb_is_leaf(cn)) {
        /* Use val as a phantom leaf to make the algorithm work with no sentinel leaf */
        if (is_black(val)) {
            do_balance(tree, val);
        }

        /* Remove the val */
        if (rb_is_root(val)) {
            tree->rn = cn;
        }
        else if (rb_is_child_left(val)) {
            rb_set_left(rb_parent(val), cn);
        }
        else {
            rb_set_right(rb_parent(val), cn);
        }
    }
    else {
        /* Remove the val */
        if (rb_is_root(val)) {
            tree->rn = cn;
        }
        else if (rb_is_child_left(val)) {
            rb_set_left(rb_parent(val), cn);
        }
        else {
            rb_set_right(rb_parent(val), cn);
        }
        rb_set_parent(cn, rb_parent(val));

        /* Re-balance */
        if (is_black(val)) {
            if (is_red(cn)) {
                set_black(cn);
            }
            else {
                do_balance(tree, cn);
            }
        }
    }

    list_del(&val->list);
}

void rb_tree_delete(struct rb_tree *tree, struct rb_tree_node *node) 
{
    if (!rb_is_leaf(rb_left(node)) && !rb_is_leaf(rb_right(node))) {
        struct rb_tree_node *child  = rb_right(node);

        while(!rb_is_leaf(rb_left(child))) {
            child = rb_left(child);
        }
     
        rb_xchg_node(tree, node, child);
    }
    update_delete(tree, node);

    rb_tree_free(tree, node);
}

struct rb_tree_node* rb_tree_find(struct rb_tree *tree, struct rb_tree_node *node) 
{
    struct rb_tree_node *rn = tree->rn;
    int cmp;

    while (rn && (cmp = rb_tree_comp(tree, rn, node))) {
        if (cmp > 0) {
            rn = rb_left(rn);
        }
        else {
            rn = rb_right(rn);
        }
    }

    return rn;
}

void rb_tree_clear(struct rb_tree *tree)
{
    struct rb_tree_node *tmp;
    struct list_head *pos, *n;

    list_for_each_safe(pos, n, &tree->fn) {
        tmp = list_entry(pos, struct rb_tree_node, list);
        list_del(pos);
        rb_tree_free(tree, tmp);
    }

    INIT_RB_TREE(tree, tree->_comp, tree->_free);
}

static int check_tree(struct rb_tree_node *node, int *res)
{
    int l, r;

    if (rb_is_leaf(node)) {
        return is_black(node);
    }
    l = is_black(node) + check_tree(rb_left(node), res);
    r = is_black(node) + check_tree(rb_right(node), res);

    if (l != r) {
        *res = FAILURE;
    }
    if (is_red(node)) {
        if (is_red(rb_left(node)) || is_red(rb_right(node))) {
            *res = FAILURE;
        }
    }

    return l;
}

int rb_tree_check(struct rb_tree *tree)
{
    int res = SUCCESS;

    if (!rb_tree_empty(tree)) {
        if (!rb_is_root(tree->rn)) {
            res = FAILURE;
        }
        else if (!is_black(tree->rn)) {
            res = FAILURE;
        }
        else {
            check_tree(tree->rn, &res);
        }
    }
    
    return res;
}
